
import numpy as np
import matplotlib.pyplot as plt
import copy
import torch
import torch.nn as nn
from nets.vgg16 import vgg16
from nets.fused_mobilenet_v3 import mobilenet_v3
# from nets.fused_mobilenet_v3_LMSoftmax import mobilenet_v3
# from nets.light_fused_mobilenet_v3 import mobilenet_v3
from nets.ghostnet import ghostnet
from nets.efficientnet import EfficientNet
from nets.shufflenet_v2 import shufflenetv2
from utils.utils import get_classes, predict_input



class Classification(object):
    _defaults = {
        # --------------------------------------------------------------------------#
        #   使用自己训练好的模型进行预测一定要修改model_path和classes_path！
        #   model_path指向logs文件夹下的权值文件，classes_path指向model_data下的txt
        #   如果出现shape不匹配，同时要注意训练时的model_path和classes_path参数的修改
        # --------------------------------------------------------------------------#
        "model_path": 'logs/loss_acc_lr_2024_03_05_19_12_13_fused2_Eca_spa1_siamese/best_epoch_weights.pth',
        # "model_path": 'logs/loss_acc_lr_2024_03_25_09_55_21_fused2_Eca_spa1_lmsoftmax/best_epoch_weights.pth',
        # "model_path": 'net_light/best_epoch_weights.pth',
        "classes_path": 'model_data/news_classes.txt',
        # --------------------------------------------------------------------#
        #   输入的图片大小
        # --------------------------------------------------------------------#
        "input_shape": [224, 224],
        # --------------------------------------------------------------------#
        #   所用模型种类：
        #   LeNet、AlexNet、vgg16、
        #   squeezenet、mobilenetv2、mobilenetv3、shufflenetv2、ghostnet、
        #   DenseNet、resnet50、resnext50
        # --------------------------------------------------------------------#
        "backbone": 'mobilenetv3-Large-Margin Softmax',
        # -------------------------------#
        #   是否使用Cuda
        #   没有GPU可以设置成False
        # -------------------------------#
        "cuda": True
    }

    # ---------------------------------------------------#
    #   初始化classification
    # ---------------------------------------------------#
    def __init__(self, **kwargs):
        self.__dict__.update(self._defaults)  # 把_defaults中的参数，初始化为类的属性
        for name, value in kwargs.items():
            setattr(self, name, value)  # 设置类的属性

        # ---------------------------------------------------#
        #   获得种类
        # ---------------------------------------------------#
        self.class_names, self.num_classes = get_classes(self.classes_path)
        self.generate()

    # ---------------------------------------------------#
    #   获得所有的分类
    # ---------------------------------------------------#
    def generate(self):
        # ---------------------------------------------------#
        #   载入模型与权值
        # ---------------------------------------------------#
        # if self.backbone == "LeNet":
        #     self.model = get_model_from_name[self.backbone](num_classes=self.num_classes, in_channels=3)
        # elif self.backbone == "mobilenetv3":
        #     self.model = get_model_from_name[self.backbone](n_class=self.num_classes, mode='large')
        # elif self.backbone == "DenseNet":
        #     self.model = get_model_from_name[self.backbone](num_classes=self.num_classes, small_inputs=False)
        # elif self.backbone == "shufflenetv2":
        #     self.model = get_model_from_name[self.backbone](class_num=self.num_classes, ratio=1.5)
        # elif self.backbone == "squeezenet":
        #     self.model = get_model_from_name[self.backbone](version=1.1, num_classes=self.num_classes)
        # else:
        #     self.model = get_model_from_name[self.backbone](num_classes=self.num_classes)
        self.model = mobilenet_v3(pretrained=False, mode='small', num_classes=self.num_classes)
        # self.model = ghostnet(num_classes=self.num_classes)
        # self.model = EfficientNet.from_name('efficientnet-b0', override_params={'num_classes': self.num_classes})
        # self.model = shufflenetv2(ratio=1, class_num=self.num_classes)
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        # self.model.load_state_dict(torch.load(self.model_path, map_location=device))
        self.model = self.model.eval()
        print('{} {} model, and classes loaded.'.format(self.backbone, self.model_path))

        if self.cuda:
            self.model = nn.DataParallel(self.model)
            self.model = self.model.cuda()

    # ---------------------------------------------------#
    #   检测图片
    # ---------------------------------------------------#
    def detect_image(self, image):

        origin_image = copy.deepcopy(image)
        # ---------------------------------------------#
        #   对测试图片预处理
        # ---------------------------------------------#
        image_data = predict_input(image, self.input_shape)
        # ---------------------------------------------------------#
        #   添加batch_size维度
        # ---------------------------------------------------------#
        image_data = np.expand_dims(image_data, 0)

        with torch.no_grad():
            photo = torch.from_numpy(image_data)
            if self.cuda:
                photo = photo.cuda()
            # ---------------------------------------------------#
            #   图片传入网络进行预测
            # ---------------------------------------------------#
            preds = torch.softmax(self.model(photo)[0], dim=-1).cpu().numpy()
        # ---------------------------------------------------#
        #   获得所属种类
        # ---------------------------------------------------#
        class_name = self.class_names[np.argmax(preds)]
        probability = np.max(preds)

        # ---------------------------------------------------#
        #   绘图并写字
        # ---------------------------------------------------#
        plt.subplot(1, 1, 1)
        plt.imshow(np.array(origin_image))
        plt.title('Class:%s Probability:%.3f' % (class_name, probability))
        plt.show()
        return class_name
